
import openai
import os
from langchain import OpenAI, SQLDatabase, SQLDatabaseChain

openai.api_key = os.environ.get("OPENAI_API_KEY")

class Conversation:
    def __init__(self):
        db = SQLDatabase.from_uri(
            "mysql+pymysql://root:password@115.159.119.116:3306/openai_demo",
            include_tables=['data_wide_s_trade'],
            sample_rows_in_table_info=2
        )
        llm = OpenAI(temperature=0)
        self.db_chain = SQLDatabaseChain.from_llm(llm=llm, db=db, verbose=True)

    def ask(self, question):
        try:
            return self.db_chain.run(question)
        except Exception as e:
            try:
                print(e)
                response = openai.Completion.create(
                    engine="text-davinci-003",
                    prompt=question,
                    max_tokens=512,
                    n=1,
                    stop=None,
                    temperature=0.5,
                )
                print(response)
                message = response.choices[0].text.strip()
                return message
            except Exception as e:
                print(e)
                return "AI也不知道"

import gradio as gr

conv = Conversation()

def answer(question, history=[]):
    history.append(question)
    response = conv.ask(question)
    history.append(response)
    responses = [(u,b) for u,b in zip(history[::2], history[1::2])]
    print(response)
    return responses, history

with gr.Blocks(css="#chatbot{height:500px} .overflow-y-auto{height:800px}") as demo:
    chatbot = gr.Chatbot(elem_id="chatbot")
    state = gr.State([])

    with gr.Row():
        txt = gr.Textbox(show_label=False, placeholder="Enter text and press enter").style(container=False)

    txt.submit(answer, [txt, state], [chatbot, state])

demo.launch()